// Copyright (c) Lawrence Livermore National Security, LLC and other VisIt
// Project developers.  See the top-level LICENSE file for dates and other
// details.  No copyright assignment is required to contribute to VisIt.

// ************************************************************************* //
//                            avtADIOS2SSTFileFormat.C                           //
// ************************************************************************* //

#include <avtMTSDFileFormatInterface.h>
#include <avtADIOS2SSTFileFormat.h>

#include <string>
#include <map>

#include <vtkFloatArray.h>
#include <vtkDoubleArray.h>
#include <vtkRectilinearGrid.h>
#include <vtkStructuredGrid.h>
#include <vtkUnstructuredGrid.h>

#include <avtDatabaseMetaData.h>

#include <Expression.h>

#include <InvalidVariableException.h>
#include <VisItStreamUtil.h>
#include <DebugStream.h>

using namespace std;

static std::vector<int> CSVToVectorInt(const std::string csv) noexcept
{
    std::vector<int> numbers;
    if (csv.empty())
    {
        return numbers;
    }

    if (csv.find(",") == csv.npos) // if no commas, one int
    {
        numbers.push_back(std::stoi(csv)); // might need to be checked
    }
    else
    {
        numbers.reserve(std::count(csv.begin(), csv.end(), ','));

        std::istringstream csvSS(csv);
        std::string value;
        while (std::getline(csvSS, value, ','))
        {
            numbers.push_back(std::stoi(csv));
        }
    }

    return numbers;
}

bool
avtADIOS2SSTFileFormat::Identify(const char *fname)
{
    return false;
}

avtFileFormatInterface *
avtADIOS2SSTFileFormat::CreateInterface(const char *const *list,
                                         int nList,
                                         int nBlock)
{
    int nTimestepGroups = nList / nBlock;
    avtMTSDFileFormat ***ffl = new avtMTSDFileFormat**[nTimestepGroups];
    for (int i = 0; i < nTimestepGroups; i++)
    {
        ffl[i] =  new avtMTSDFileFormat*[nBlock];
        for (int j = 0; j < nBlock; j++)
            ffl[i][j] =  new avtADIOS2SSTFileFormat(list[i*nBlock +j]);
    }
    return new avtMTSDFileFormatInterface(ffl, nTimestepGroups, nBlock);
}

// ****************************************************************************
//  Method: avtADIOS2SSTFileFormat constructor
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Thu Apr 12 08:39:47 PDT 2018
// 
//  Modifications:
//    Justin Privitera, Thu Jan 18 09:56:51 PST 2024
//    Removed adios2::DebugON since it is not present in newer versions of 
//    adios2.
//
// ****************************************************************************

avtADIOS2SSTFileFormat::avtADIOS2SSTFileFormat(const char *filename)
    : adios(std::make_shared<adios2::ADIOS>()),
      io(adios->DeclareIO("sstIO")),
      numTimeSteps(1),
      avtMTSDFileFormat(&filename, 1)
{
    io.SetEngine("SST");
    debug5<<"Open "<<filename<<endl;

    debug5<<__FILE__<<" "<<__LINE__<<endl;
    reader = io.Open(filename, adios2::Mode::Read);
    debug5<<__FILE__<<" "<<__LINE__<<endl;
    variables = io.AvailableVariables();
    debug5<<__FILE__<<" "<<__LINE__<<endl;

    if (variables.size() > 0)
    {
        auto var0 = variables.begin()->second;
        string nsteps = var0["AvailableStepsCount"];
        numTimeSteps = std::stoi(nsteps);
    }

    for (auto &v : variables)
        debug5<<"Var: "<<v.first<<endl;
}


// ****************************************************************************
//  Method: avtADIOS2SSTFileFormat::GetNTimesteps
//
//  Purpose:
//      Tells the rest of the code how many timesteps there are in this file.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Thu Apr 12 08:39:47 PDT 2018
//
// ****************************************************************************

int
avtADIOS2SSTFileFormat::GetNTimesteps()
{
    debug5<<"avtADIOS2SSTFileFormat::GetNTimesteps()"<<endl;
    return numTimeSteps;
}


// ****************************************************************************
//  Method: avtADIOS2FileFormat::FreeUpResources
//
//  Purpose:
//      When VisIt is done focusing on a particular timestep, it asks that
//      timestep to free up any resources (memory, file descriptors) that
//      it has associated with it.  This method is the mechanism for doing
//      that.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Thu Apr 12 08:39:47 PDT 2018
//
// ****************************************************************************

void
avtADIOS2SSTFileFormat::FreeUpResources(void)
{
}


// ****************************************************************************
//  Method: avtADIOS2SSTFileFormat::PopulateDatabaseMetaData
//
//  Purpose:
//      This database meta-data object is like a table of contents for the
//      file.  By populating it, you are telling the rest of VisIt what
//      information it can request from you.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Thu Apr 12 08:39:47 PDT 2018
//
// ****************************************************************************

void
avtADIOS2SSTFileFormat::PopulateDatabaseMetaData(avtDatabaseMetaData *md, int timeState)
{
    vector<pair<string,string>> vars;

    for (const auto &varInfo : variables)
    {
        auto params = varInfo.second;
        string shape = params["Shape"];
        auto vecShape = CSVToVectorInt(shape);
        if (meshInfo.find(shape) == meshInfo.end())
        {
            string meshNm;
            int dim = 0;
            if (vecShape.size() == 2)
            {
                meshNm = "mesh" + std::to_string(vecShape[0])+"x"+std::to_string(vecShape[1]);
                dim = 2;
            }
            else if (vecShape.size() == 3)
            {
                meshNm = "mesh" + std::to_string(vecShape[0])+"x"+std::to_string(vecShape[1])+"x"+std::to_string(vecShape[2]);
                dim = 3;
            }
            meshInfo[shape] = std::make_pair(dim, meshNm);
        }
        vars.push_back(std::make_pair(varInfo.first, meshInfo[shape].second));
    }

    for (auto &m : meshInfo)
    {
        string nm = m.second.second;
        int dim = m.second.first;
        AddMeshToMetaData(md, nm, AVT_RECTILINEAR_MESH, NULL, 1, 0, dim, dim);
    }

    for (auto &v : vars)
        AddScalarVarToMetaData(md, v.first, v.second, AVT_NODECENT);
}


// ****************************************************************************
//  Method: avtADIOS2SSTFileFormat::GetMesh
//
//  Purpose:
//      Gets the mesh associated with this file.  The mesh is returned as a
//      derived type of vtkDataSet (ie vtkRectilinearGrid, vtkStructuredGrid,
//      vtkUnstructuredGrid, etc).
//
//  Arguments:
//      timestate   The index of the timestate.  If GetNTimesteps returned
//                  'N' time steps, this is guaranteed to be between 0 and N-1.
//      meshname    The name of the mesh of interest.  This can be ignored if
//                  there is only one mesh.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Thu Apr 12 08:39:47 PDT 2018
//
// ****************************************************************************

vtkDataSet *
avtADIOS2SSTFileFormat::GetMesh(int timestate, const char *meshname)
{
    //Find the mesh info we need.
    string shapeStr;
    for (auto &m : meshInfo)
    {
        if (m.second.second == string(meshname))
        {
            shapeStr = m.first;
            break;
        }
    }
    if (shapeStr.size() == 0)
        return NULL;

    auto shape = CSVToVectorInt(shapeStr);
    if (shape.size() == 2)
        shape.push_back(1);

    int dims[3] = {shape[0], shape[1], shape[2]};

    vtkRectilinearGrid *grid = vtkRectilinearGrid::New();
    grid->SetDimensions(dims);
    vtkFloatArray *coords[3] = {vtkFloatArray::New(),vtkFloatArray::New(),vtkFloatArray::New()};

    for (int c = 0; c < 3; c++)
    {
        coords[c]->SetNumberOfTuples(shape[c]);
        for (int i = 0; i < shape[c]; i++)
            coords[c]->SetTuple1(i, i);
    }

    grid->SetXCoordinates(coords[0]);
    grid->SetYCoordinates(coords[1]);
    grid->SetZCoordinates(coords[2]);
    coords[0]->Delete();
    coords[1]->Delete();
    coords[2]->Delete();

    return grid;
}


// ****************************************************************************
//  Method: avtADIOS2SSTFileFormat::GetVar
//
//  Purpose:
//      Gets a scalar variable associated with this file.  Although VTK has
//      support for many different types, the best bet is vtkFloatArray, since
//      that is supported everywhere through VisIt.
//
//  Arguments:
//      timestate  The index of the timestate.  If GetNTimesteps returned
//                 'N' time steps, this is guaranteed to be between 0 and N-1.
//      varname    The name of the variable requested.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Thu Apr 12 08:39:47 PDT 2018
//
// ****************************************************************************

vtkDataArray *
avtADIOS2SSTFileFormat::GetVar(int timestate, const char *varname)
{
    debug5<<"GetVar: "<<varname<<endl;

    if (variables.find(varname) == variables.end())
        return NULL;

    auto var = variables[varname];
    string varType = var["Type"];
    if (varType == "double")
    {
        adios2::Variable<double> v = io.InquireVariable<double>(varname);
        debug5<<"DIMS= "<<v.Shape()<<endl;
        v.SetSelection(adios2::Box<adios2::Dims>({0,0,0}, v.Shape()));
        v.SetStepSelection({timestate, 1});

        size_t numVals = 1;
        for (int i = 0; i < v.Shape().size(); i++)
            numVals *= v.Shape()[i];

        vector<double> buff(numVals);
        reader.Get(v, buff.data(), adios2::Mode::Sync);

        vtkDoubleArray *arr = vtkDoubleArray::New();
        arr->SetNumberOfComponents(1);
        arr->SetNumberOfTuples(numVals);
        for (int i = 0; i < numVals; i++)
            arr->SetTuple1(i, buff[i]);
        return arr;
    }

    return NULL;
}


// ****************************************************************************
//  Method: avtADIOS2SSTFileFormat::GetVectorVar
//
//  Purpose:
//      Gets a vector variable associated with this file.  Although VTK has
//      support for many different types, the best bet is vtkFloatArray, since
//      that is supported everywhere through VisIt.
//
//  Arguments:
//      timestate  The index of the timestate.  If GetNTimesteps returned
//                 'N' time steps, this is guaranteed to be between 0 and N-1.
//      varname    The name of the variable requested.
//
//  Programmer: pugmire -- generated by xml2avt
//  Creation:   Thu Apr 12 08:39:47 PDT 2018
//
// ****************************************************************************

vtkDataArray *
avtADIOS2SSTFileFormat::GetVectorVar(int timestate, const char *varname)
{
    return NULL;
}
